{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e9b67513",
   "metadata": {},
   "source": [
    "# Deep Dive into Molecular Search Engine with Towhee"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fd00bbc5",
   "metadata": {},
   "source": [
    "In the [previous tutorial](1_build_molecular_search_engine.ipynb), we built and prototyped a proof-of-concept molecular similar search engine. Now, let's try substructure and superstructure search, and deploy it as a micro-service with Towhee."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6d85d5be",
   "metadata": {},
   "source": [
    "## Preparation\n",
    "If you haven't done so already, please go through our previous tutorial: \"[Build a Molecular Search Engine in Minutes](./1_build_molecular_search_engine.ipynb)\". To make things easy, we'll repeat the important code blocks below; if you have already executed these blocks, please move on to next section."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6eef0af6",
   "metadata": {},
   "source": [
    "### Install Dependencies\n",
    "First we need to install dependencies such as towhee, rdkit and gradio."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "b19c410a",
   "metadata": {},
   "outputs": [],
   "source": [
    "! python -m pip install -q towhee towhee.models rdkit-pypi gradio"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e2a6cdcd",
   "metadata": {},
   "source": [
    "Please make sure that you have started a [Milvus service](https://milvus.io/docs/install_standalone-docker.md). This notebook uses [milvus 2.2.10](https://milvus.io/docs/v2.2.x/install_standalone-docker.md) and [pymilvus 2.2.11](https://milvus.io/docs/release_notes.md#2210)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "7ab3de54",
   "metadata": {},
   "outputs": [],
   "source": [
    "! python -m pip install -q pymilvus==2.2.11"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "135039d7",
   "metadata": {},
   "source": [
    "### Prepare the Data\n",
    "There is a subset of the [Pubchem dataset](https://ftp.ncbi.nlm.nih.gov/pubchem/Compound/CURRENT-Full/SDF/) (10000 SMILES) used in this demo, everyone can download on [Github](https://github.com/towhee-io/examples/releases/download/data/pubchem_10000.smi)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "89df4afd",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  % Total    % Received % Xferd  Average Speed   Time    Time     Time  Current\n",
      "                                 Dload  Upload   Total   Spent    Left  Speed\n",
      "  0     0    0     0    0     0      0      0 --:--:-- --:--:-- --:--:--     0\n",
      "100  561k  100  561k    0     0   242k      0  0:00:02  0:00:02 --:--:--  438k\n"
     ]
    }
   ],
   "source": [
    "! curl -L https://github.com/towhee-io/examples/releases/download/data/pubchem_10000.smi -O"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a90cc55",
   "metadata": {},
   "source": [
    "To use the dataset for molecular search, let's first define the dictionary and helper function:\n",
    "- `to_images(input)`: convert the input smiles or results to towhee.Image for display."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "c9728c01",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "from rdkit.Chem import Draw\n",
    "from rdkit import Chem\n",
    "from towhee.types.image_utils import from_pil\n",
    "\n",
    "df = pd.read_csv('pubchem_10000.smi')\n",
    "\n",
    "def to_images(data):\n",
    "    imgs = []\n",
    "    for smiles in data:\n",
    "        mol = Chem.MolFromSmiles(smiles)\n",
    "        img = from_pil(Draw.MolToImage(mol))\n",
    "        imgs.append(img)\n",
    "    return imgs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "69b5b4eb",
   "metadata": {},
   "source": [
    "### Create Milvus Collection"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5e1d18ee",
   "metadata": {},
   "source": [
    "Next to create two collection for substructure and substru\"molsearch\" collection in Milvus. It's worth to note that the vector field of this collection is the type of \"BINARY_VECTOR\"."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "68af7917",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pymilvus import connections, FieldSchema, CollectionSchema, DataType, Collection, utility\n",
    "\n",
    "MILVUS_HOST = '127.0.0.1'\n",
    "MILVUS_PORT = '19530'\n",
    "connections.connect(host=MILVUS_HOST, port=MILVUS_PORT)\n",
    "\n",
    "def create_milvus_collection(collection_name, dim):    \n",
    "    if utility.has_collection(collection_name):\n",
    "        utility.drop_collection(collection_name)\n",
    "    \n",
    "    fields = [\n",
    "    FieldSchema(name='id', dtype=DataType.INT64, descrition='ids', is_primary=True, auto_id=False),\n",
    "    FieldSchema(name='smiles', dtype=DataType.VARCHAR, descrition='SMILES', max_length=500),\n",
    "    FieldSchema(name='embedding', dtype=DataType.BINARY_VECTOR, descrition='embedding vectors', dim=dim)\n",
    "    ]\n",
    "    schema = CollectionSchema(fields=fields, description='molecular similarity search')\n",
    "    collection = Collection(name=collection_name, schema=schema)\n",
    "    \n",
    "    index_params = {\"index_type\": \"BIN_FLAT\", \"params\": {\"nlist\": 1024}, \"metric_type\": \"JACCARD\"}\n",
    "    collection.create_index(field_name=\"embedding\", index_params=index_params)\n",
    "    \n",
    "    return collection\n",
    "\n",
    "collection = create_milvus_collection('molecular_search', 2048)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "78b1a0a3",
   "metadata": {},
   "source": [
    "Then load smiles data to Milvus with Towhee, more details in [\"Build a Molecular Search Engine in Minutes\"](./1_build_molecular_search_engine.ipynb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "29e2fefa",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total number of inserted data is 10000.\n"
     ]
    }
   ],
   "source": [
    "from towhee import ops, pipe\n",
    "\n",
    "insert_pipe = (pipe.input('df')\n",
    "                   .flat_map('df', ('smiles', 'id'), lambda df: df.values.tolist())\n",
    "                   .map('smiles', 'fp', ops.molecular_fingerprinting.rdkit(algorithm='daylight'))\n",
    "                   .map(('id', 'smiles', 'fp'), 'res', ops.ann_insert.milvus_client(host=MILVUS_HOST, \n",
    "                                                                           port=MILVUS_PORT,\n",
    "                                                                           collection_name='molecular_search'))\n",
    "                   .output('res')\n",
    ")\n",
    "\n",
    "insert_pipe(df)\n",
    "\n",
    "collection.flush()\n",
    "print('Total number of inserted data is {}.'.format(collection.num_entities))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e6c79306",
   "metadata": {},
   "source": [
    "## Sub and Super Structure Search\n",
    "\n",
    "Milvus not only supports searching similar structures of molecular formulas, but also superstructure and substructure searches, you only need to specify the metric types:\n",
    "\n",
    "- Similarly search: \"JACCARD\"\n",
    "- Superstructure search: \"SUPERSTRUCTURE\"\n",
    "- Substructure search: \"SUBSTRUCTURE\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "0c0dc6d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "search_pipe = (pipe.input('query_smiles')\n",
    "                   .map('query_smiles', 'fp', ops.molecular_fingerprinting.rdkit(algorithm='daylight'))\n",
    "                   .flat_map('fp', ('super_id', 'super_score', 'super_smiles'), \n",
    "                                     ops.ann_search.milvus_client(host=MILVUS_HOST, \n",
    "                                                                  port=MILVUS_PORT,\n",
    "                                                                  collection_name='molecular_search',\n",
    "                                                                  limit=6,\n",
    "                                                                  param={\"metric_type\": \"SUPERSTRUCTURE\", \"nprobe\": 10},\n",
    "                                                                  output_fields=['smiles']))\n",
    "                   .flat_map('fp', ('sub_id', 'sub_score', 'sub_smiles'), \n",
    "                                     ops.ann_search.milvus_client(host=MILVUS_HOST, \n",
    "                                                                  port=MILVUS_PORT,\n",
    "                                                                  collection_name='molecular_search',\n",
    "                                                                  limit=6,\n",
    "                                                                  param={\"metric_type\": \"SUBSTRUCTURE\", \"nprobe\": 10},\n",
    "                                                                  output_fields=['smiles']))\n",
    "                   .window_all('query_smiles', 'query_smiles', lambda x: to_images(x[:1]))\n",
    "                   .window_all('sub_smiles', 'sub_smiles', to_images)\n",
    "                   .window_all('super_smiles', 'super_smiles', to_images)\n",
    "                   .output('query_smiles', 'sub_smiles', 'super_smiles')\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "43f48ff8",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/chenshiyu/workspace/git/towhee/towhee/datacollection/mixins/display.py:215: RuntimeWarning: More than 20 figures have been opened. Figures created through the pyplot interface (`matplotlib.pyplot.figure`) are retained until explicitly closed and may consume too much memory. (To control this warning, see the rcParam `figure.max_open_warning`).\n",
      "  fig = plt.figure(figsize=(width / 100, height / 100))\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<table style=\"border-collapse: collapse;\"><tr><th style=\"text-align: center; font-size: 130%; border: none;\">query_smiles</th> <th style=\"text-align: center; font-size: 130%; border: none;\">sub_smiles</th> <th style=\"text-align: center; font-size: 130%; border: none;\">super_smiles</th></tr> <tr><td style=\"text-align: center; vertical-align: top; border-right: solid 1px #D3D3D3; border-left: solid 1px #D3D3D3; \"><img src=\"\" 128 = \"128px\" 128 = \"128px\" style = \"float:left; padding:2px\"></td> <td style=\"text-align: center; vertical-align: top; border-right: solid 1px #D3D3D3; border-left: solid 1px #D3D3D3; \"><img src=\"\" 128 = \"128px\" 128 = \"128px\" style = \"float:left; padding:2px\"> <img src=\"\" 128 = \"128px\" 128 = \"128px\" style = \"float:left; padding:2px\"> <img src=\"\" 128 = \"128px\" 128 = \"128px\" style = \"float:left; padding:2px\"> <img src=\"\" 128 = \"128px\" 128 = \"128px\" style = \"float:left; padding:2px\"> <img src=\"\" 128 = \"128px\" 128 = \"128px\" style = \"float:left; padding:2px\"> <img src=\"\" 128 = \"128px\" 128 = \"128px\" style = \"float:left; padding:2px\"> <img src=\"\" 128 = \"128px\" 128 = \"128px\" style = \"float:left; padding:2px\"> <img src=\"\" 128 = \"128px\" 128 = \"128px\" style = \"float:left; padding:2px\"> <img src=\"\" 128 = \"128px\" 128 = \"128px\" style = \"float:left; padding:2px\"> <img src=\"\" 128 = \"128px\" 128 = \"128px\" style = \"float:left; padding:2px\"> <img src=\"\" 128 = \"128px\" 128 = \"128px\" style = \"float:left; padding:2px\"> <img src=\"\" 128 = \"128px\" 128 = \"128px\" style = \"float:left; padding:2px\"> <img src=\"\" 128 = \"128px\" 128 = \"128px\" style = \"float:left; padding:2px\"> <img src=\"\" 128 = \"128px\" 128 = \"128px\" style = \"float:left; padding:2px\"> <img src=\"\" 128 = \"128px\" 128 = \"128px\" style = \"float:left; padding:2px\"> <img src=\"\" 128 = \"128px\" 128 = \"128px\" style = \"float:left; padding:2px\"> <img src=\"\" 128 = \"128px\" 128 = \"128px\" style = \"float:left; padding:2px\"> <img src=\"\" 128 = \"128px\" 128 = \"128px\" style = \"float:left; padding:2px\"></td> <td style=\"text-align: center; vertical-align: top; border-right: solid 1px #D3D3D3; border-left: solid 1px #D3D3D3; \"><img src=\"\" 128 = \"128px\" 128 = \"128px\" style = \"float:left; padding:2px\"> <img src=\"\" 128 = \"128px\" 128 = \"128px\" style = \"float:left; padding:2px\"> <img src=\"\" 128 = \"128px\" 128 = \"128px\" style = \"float:left; padding:2px\"></td></tr></table>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from towhee import DataCollection\n",
    "\n",
    "res = search_pipe('Cn1ccc(=O)nc1')\n",
    "DataCollection(res).show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5b5addf5",
   "metadata": {},
   "source": [
    "## Release a Showcase\n",
    "\n",
    "We've done an excellent job on the core functionality of our molecular search engine. Now it's time to build a showcase with interface. Gradio is a great tool for building demos. With Gradio, we simply need to wrap the data processing pipeline via a `search_smiles_with_metric` function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "12c11744",
   "metadata": {},
   "outputs": [],
   "source": [
    "def search_smiles_with_metric(smiles, metric_type):\n",
    "    search_func = (pipe.input('query_smiles')\n",
    "                   .map('query_smiles', 'fp', ops.molecular_fingerprinting.rdkit(algorithm='daylight'))\n",
    "                   .flat_map('fp', ('id', 'score', 'similar_smiles'), ops.ann_search.milvus_client(host=MILVUS_HOST, \n",
    "                                                                  port=MILVUS_PORT,\n",
    "                                                                  collection_name='molecular_search',\n",
    "                                                                  limit=5,\n",
    "                                                                  param={\"metric_type\": metric_type, \"nprobe\": 10},\n",
    "                                                                  output_fields=['smiles']))\n",
    "                   .window_all('query_smiles', 'query_smiles', lambda x: to_images(x[:1]))\n",
    "                   .window_all('similar_smiles', 'similar_smiles', to_images)\n",
    "                   .output('similar_smiles')\n",
    "    )\n",
    "    return search_func(smiles).to_list()[0][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "e0ffb773",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/chenshiyu/miniconda3/envs/dev-towhee/lib/python3.8/site-packages/gradio/inputs.py:26: UserWarning: Usage of gradio.inputs is deprecated, and will not be supported in the future, please import your component from gradio.components\n",
      "  warnings.warn(\n",
      "/Users/chenshiyu/miniconda3/envs/dev-towhee/lib/python3.8/site-packages/gradio/deprecation.py:40: UserWarning: `optional` parameter is deprecated, and it has no effect\n",
      "  warnings.warn(value)\n",
      "/Users/chenshiyu/miniconda3/envs/dev-towhee/lib/python3.8/site-packages/gradio/deprecation.py:40: UserWarning: `numeric` parameter is deprecated, and it has no effect\n",
      "  warnings.warn(value)\n",
      "/Users/chenshiyu/miniconda3/envs/dev-towhee/lib/python3.8/site-packages/gradio/deprecation.py:40: UserWarning: The 'type' parameter has been deprecated. Use the Number component instead.\n",
      "  warnings.warn(value)\n",
      "/Users/chenshiyu/miniconda3/envs/dev-towhee/lib/python3.8/site-packages/gradio/inputs.py:182: UserWarning: Usage of gradio.inputs is deprecated, and will not be supported in the future, please import your component from gradio.components\n",
      "  warnings.warn(\n",
      "/Users/chenshiyu/miniconda3/envs/dev-towhee/lib/python3.8/site-packages/gradio/outputs.py:42: UserWarning: Usage of gradio.outputs is deprecated, and will not be supported in the future, please import your components from gradio.components\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "IMPORTANT: You are using gradio version 3.3, however version 3.14.0 is available, please upgrade.\n",
      "--------\n",
      "Running on local URL:  http://127.0.0.1:7861\n",
      "Running on public URL: https://9d6bc9cc75628221.gradio.app\n",
      "\n",
      "This share link expires in 72 hours. For free permanent hosting, check out Spaces: https://www.huggingface.co/spaces\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div><iframe src=\"https://9d6bc9cc75628221.gradio.app\" width=\"900\" height=\"500\" allow=\"autoplay; camera; microphone;\" frameborder=\"0\" allowfullscreen></iframe></div>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "(<gradio.routes.App at 0x7f97e2c7dca0>,\n",
       " 'http://127.0.0.1:7861/',\n",
       " 'https://9d6bc9cc75628221.gradio.app')"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import gradio\n",
    "\n",
    "interface = gradio.Interface(search_smiles_with_metric, \n",
    "                             [gradio.inputs.Textbox(lines=1, default='CN1C=NC2=C1C(=O)N(C(=O)N2C)C'), \n",
    "                              gradio.inputs.Radio(['JACCARD', 'SUBSTRUCTURE', 'SUPERSTRUCTURE'])],\n",
    "                             [gradio.outputs.Image(type=\"pil\", label=None) for _ in range(5)]\n",
    "                            )\n",
    "\n",
    "interface.launch(inline=True, share=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2673be6a",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
